import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle 
from matplotlib.animation import ImageMagickBase
import torch
from typing import List
import numpy as np
import PIL.Image as Image

def draw_image_with_grad(image_with_grad,grad=None):
    if isinstance(image_with_grad,List):
        image = [i.detach().cpu().numpy().transpose(1,2,0) for i in image_with_grad] #B,W,H,chanels
        if not grad:
            grad = [i.grad.data.abs() for i in image_with_grad]
        saliency = [torch.max(g,dim=0)[0] for g in grad]
        saliency = torch.stack(saliency,dim=0)# B,W,H
    elif isinstance(image_with_grad,torch.Tensor):
        image= image_with_grad.detach().cpu().numpy().transpose(0,2,3,1) # B,W,H,3
        if not grad:
            grad = image_with_grad.grad.data.abs() # B,C,W,H
        saliency = torch.max(grad, dim=1)[0]# B,W,H
    #print(saliency.shape)
    saliency = saliency.unsqueeze(-1).detach().cpu().numpy() # B,W,H,1
    saliency = (saliency- saliency.min(axis=0))/(saliency.max(axis=0)-saliency.min(axis=0)+1e-8)
    saliency = np.concatenate([saliency,np.zeros_like(saliency),np.zeros_like(saliency)],axis=-1) # B,W,H,3
    return image, saliency

def fig2data(fig):
    """
    fig = plt.figure()
    image = fig2data(fig)
    @brief Convert a Matplotlib figure to a 4D numpy array with RGBA channels and return it
    @param fig a matplotlib figure
    @return a numpy 3D array of RGBA values
    """
    # draw the renderer
    fig.canvas.draw()
 
    # Get the RGBA buffer from the figure
    w, h = fig.canvas.get_width_height()
    buf = np.fromstring(fig.canvas.tostring_argb(), dtype=np.uint8)
    buf.shape = (w, h, 4)
 
    # canvas.tostring_argb give pixmap in ARGB mode. Roll the ALPHA channel to have it in RGBA mode
    buf = np.roll(buf, 3, axis=2)
    image = Image.frombytes("RGBA", (w, h), buf.tostring())
    image = np.asarray(image)
    return image[:,:,:3]

def draw_average_saliency(dict_saliency_of_room,map_key_to_color,\
    figsize=(20,1),\
    title="Average saliency for type of nodes"):
    saliency_nodes_mean = {k:np.array(v).mean() for k,v in dict_saliency_of_room.items() if len(v)>0 }
    plt.figure(figsize=figsize)
    plt.title(title,fontsize=20)
    plt.imshow(np.array(list(saliency_nodes_mean.values())).reshape(1,-1),cmap="gray")
    plt.xticks(ticks=range(len(saliency_nodes_mean)),\
        labels=list(saliency_nodes_mean.keys()),\
        rotation=30)
    plt.yticks([])
    for i in range(len(saliency_nodes_mean)):
        #plt.gca().get_xticklabels()[i].set_color([c/255. for c in CMAP[i]])
        #plt.gca().get_xticklabels()[i].set_backgroundcolor("black")
        plt.text(-0.42+i,-0.42,"%.2e"%list(saliency_nodes_mean.values())[i],\
            backgroundcolor=(1,1,1,.5),\
            fontsize=8,\
            horizontalalignment="left",verticalalignment="top")

        plt.gcf().gca().add_patch(Rectangle((-0.5+i,-0.5), 
                                0.85, 0.35, 
                                fc =map_key_to_color(list(saliency_nodes_mean.keys())[i]),#[c/255. for c in CMAP[i]],  
                                ec ="black", 
                                lw = 1) ) 

def draw_saliency_distribution(dict_saliency_of_room,map_key_to_color,\
    figsize=(10,6),\
    title="saliency distribution for room types",\
    xlabel = "Saliency (axis is scaled logarithmically)",\
    ylabel = "Type of room"):
    plt.figure(figsize=figsize)
    plt.title(title,fontsize=20)
    dict_saliency_of_room_filtered = {k:np.array(v) for k,v in dict_saliency_of_room.items() if len(v)>0 }
    key_sorted = sorted(dict_saliency_of_room_filtered,key=lambda k:dict_saliency_of_room_filtered[k].mean())
    dict_saliency_of_room_filtered_sorted = {k:dict_saliency_of_room_filtered[k] for k in key_sorted}
    parts = plt.violinplot([np.log(v+1e-8) for v in dict_saliency_of_room_filtered_sorted.values()],\
        showmedians=False,\
        showextrema=False,\
        vert=False)
    for i,pc in enumerate(parts['bodies']):
        pc.set_facecolor(map_key_to_color(list(dict_saliency_of_room_filtered_sorted.keys())[i]))
        pc.set_edgecolor('black')
        pc.set_alpha(1)
    inds = 1
    for k,v in dict_saliency_of_room_filtered_sorted.items():
        quartile1, medians, quartile3 = np.percentile(v, [25, 50, 75], axis=0)
        plt.scatter(np.log(medians),inds, marker='o', color='white', s=30, zorder=3)
        plt.hlines(inds,np.log(quartile1), np.log(quartile3), color='k', linestyle='-', lw=9,alpha=0.5)
        plt.hlines(inds,np.log(v.min()+1e-8), np.log(v.max()), color='k', linestyle='-', lw=1)
        inds+=1
    plt.yticks(ticks=range(1,len(dict_saliency_of_room_filtered_sorted)+1),\
        labels=list(dict_saliency_of_room_filtered_sorted.keys()))
    plt.ylabel(ylabel)
    xmin,xmax = float("inf"), float("-inf")
    for k,v in dict_saliency_of_room_filtered_sorted.items():
        xmin = min(xmin,v.min())
        xmax = max(xmax,v.max())
    xmin_ln,xmax_ln = np.log(xmin+1e-8),np.log(xmax)
    xmin_ln_floor,xmax_ln_ceil = np.floor(xmin_ln),np.ceil(xmax_ln)
    xticks = np.linspace(xmin_ln_floor,xmax_ln_ceil,9)
    plt.xticks(ticks=xticks,labels=["%.2e"%n for n in np.exp(xticks)])
    plt.xlabel(xlabel)